#! /bin/bash

index=  proc=  yhindex=  fname=  fastq=  recursive=  zip=  

usage="$(basename "$0") -- map single-end data using bowtie2 and yaha
  -h  show help and exit
  -x  path to bowtie2 index
  -y  path to yaha index
  -p  number of cores to use
  -n  sample name
  -q  fastq file containing reads
  -r  recursive (optional)
  -z  gzip fastq files (optional)"

if [ $# -eq 0 ]; then
  echo "$usage"
  exit 1
fi

while getopts :x:p:y:n:q:rzh opt; do
  case $opt in
    h)  echo "$usage"
        exit
        ;;
    x)  index=$OPTARG
        ;;
    p)  proc=$OPTARG
        ;;
    y)  yhindex=$OPTARG
        ;;
    n)  fname=$OPTARG
        ;;
    q)  fastq=$OPTARG
        ;;
    r)  recursive=true
        ;;
    z)  zip=true
        ;;
    :)  printf "missing argument for -%s\n" "$OPTARG" >&2
        echo "$usage" >&2
        exit 1
        ;;
    \?) printf "illegal option: -%s\n" "$OPTARG" >&2
        echo "$usage" >&2
        exit 1
        ;;
  esac
done
shift $((OPTIND - 1))

map () {
  echo "Mapping ${fname}"
  bowtie2 --local --dovetail -p$proc --fr -q -R5 -N1 -x $index -U $fastq \
  | samtools view -b - > "${fname}_unsort.bam"

  # get the header
  samtools view -H "${fname}_unsort.bam" > header

  # extract clipped reads
  samtools view "${fname}_unsort.bam" \
  | awk '$6 ~ /H|S/{print}' \
  | cat header - \
  | samtools bam2fq -n - > "${fname}.umap.fastq"

  # get the unmapped reads
  samtools view -b -f 4 "${fname}_unsort.bam" \
  | samtools bam2fq -n - >> "${fname}.umap.fastq"

  rm header

  yaha -t $proc -x $yhindex -q "${fname}.umap.fastq" -L 11 -H 2000 -M 15 -osh stdout \
  | samblaster -s "${fname}.split.sam" > /dev/null

  echo "Converting to bam"
  samtools view -Sb@ $proc "${fname}.split.sam" > "${fname}.split.bam"
  rm "${fname}.split.sam"

  echo "Sorting alignment"
  samtools sort -@ $proc "${fname}_unsort.bam" > "${fname}.bam"
  rm "${fname}_unsort.bam"

  echo "Indexing alignment"
  samtools index "${fname}.bam"

  if [ "$zip" == true ]; then
    echo "Zipping fastq files"
    gzip *.fastq
  fi
}

if [ "$recursive" == true ]; then
  for directory in ./*; do
      if [ -d "$directory" ]; then
          cd $directory
          for fastq in $(ls -d *.fastq*);do
            fname=(${fastq//.fastq*/ })
            map
          done
          cd ..
      fi
  done

else
  map
fi
